package xcrypto

import (
	"bytes"
	"crypto/aes"
	"crypto/rand"

	"gitee.com/shuohe/xutils/xbuf"
)

/**
 * AES 密钥对象。
 */
type AesKey [16]byte

/**
 * @implements IKey::Algo
 */
func (this *AesKey) Algo() string {
	return "AES-128"
}

/**
 * @implements IKey::Type
 */
func (this *AesKey) Type() string {
	return KeyType_Secret
}

/**
 * @implements IKey::ExportToXBuf
 */
func (this *AesKey) ExportToXBuf(xb *xbuf.Buffer) {
	xb.PutVStr(this.Algo())
	xb.PutVStr(this.Type())

	xb.PutBytes(this[:])
}

/**
 * @implements IKey::Export
 */
func (this *AesKey) Export() []byte {
	xb := xbuf.New(len(this) + 32)
	this.ExportToXBuf(xb)

	return xb.Bytes()
}

/**
 * @implements IKey::ExportTo
 */
func (this *AesKey) ExportTo(buf []byte) int {
	xb := xbuf.WrapForWrite(buf, false)
	this.ExportToXBuf(xb)

	return xb.Wi
}

/**
 * @implements IKey.Equals
 */
func (this *AesKey) Equals(other IKey) bool {
	otherkey, ok := other.(*AesKey)
	if !ok {
		return false
	}

	return bytes.Compare(this[:], otherkey[:]) == 0
}

/**
 * @implements ISecretKey.EncryptToXBuf
 */
func (this *AesKey) EncryptToXBuf(data []byte, xb *xbuf.Buffer) {
	datalen := len(data)
	seclen := (datalen/16 + 1) * 16

	xb.EnsureWrite(seclen)
	cipher, _ := aes.NewCipher(this[:])

	i := 0
	for ; datalen-i >= 16; i += 16 {
		cipher.Encrypt(xb.Buf[xb.Wi:], data[i:])
		xb.Wi += 16
	}

	var padding [16]byte

	pbyte := byte(seclen - datalen)
	copy(padding[:], data[i:])
	for i = datalen - i; i < 16; i++ {
		padding[i] = pbyte
	}
	cipher.Encrypt(xb.Buf[xb.Wi:], padding[:])
	xb.Wi += 16
}

/**
 * @implements ISecretKey::Encrypt
 */
func (this *AesKey) Encrypt(data []byte) []byte {
	xb := xbuf.New(256)
	this.EncryptToXBuf(data, xb)

	return xb.Bytes()
}

/**
 * @implements ISecretKey.Decrypt
 */
func (this *AesKey) Decrypt(secret []byte) ([]byte, error) {
	xb := xbuf.New(256)
	err := this.DecryptToXBuf(secret, xb)
	if err != nil {
		return nil, err
	}

	return xb.Bytes(), nil
}

/**
 * @implements ISecretKey.DecryptX
 */
func (this *AesKey) DecryptToXBuf(secret []byte, xb *xbuf.Buffer) error {
	seclen := len(secret)
	if seclen%16 != 0 {
		return errSecretData()
	}

	xb.EnsureWrite(seclen)
	cipher, _ := aes.NewCipher(this[:])

	wi := xb.Wi

	for i := 0; i < seclen; i += 16 {
		cipher.Decrypt(xb.Buf[xb.Wi:], secret[i:])
		xb.Wi += 16
	}

	pbyte := xb.Buf[xb.Wi-1]
	if pbyte > 16 {
		xb.Wi = wi
		return errSecretData()
	}
	for i := 1; i < int(pbyte); i++ {
		if xb.Buf[xb.Wi-1-i] != pbyte {
			xb.Wi = wi
			return errSecretData()
		}
	}

	xb.Wi -= int(pbyte)

	return nil
}

/**
 * 根据指定类型创建相应密钥对象。
 */
func AesKeyFactory(keyType string, keyData []byte) (IKey, int, error) {
	if keyType != KeyType_Secret {
		return nil, 0, errKeyType(keyType)
	}
	if len(keyData) < 16 {
		return nil, 0, errKeyData()
	}

	k := &AesKey{}
	copy(k[:], keyData[:16])

	return k, 16, nil
}

/**
 * 生成随机数密钥。
 */
func AesKeyGenerator() (IKey, error) {
	k := &AesKey{}
	rand.Read(k[:])

	return k, nil
}
